from pyspark.mllib.evaluation import RegressionMetrics
from pyspark.mllib.recommendation import ALS, Rating, MatrixFactorizationModel
from pyspark import SparkContext
import os

os.environ['PYSPARK_PYTHON'] = "python3"

spark = SparkContext.getOrCreate()
rdd = spark.textFile("file:///Users/sonto/Workspace/P1905/spark_example/ml_lesson/ml-100k/u.data")


def map_line(line):
    data = line.split("\t")
    # return Rating(int(data[0]), int(data[1]), float(data[2]))
    return (int(data[0]), int(data[1]), float(data[2]))


rdd = rdd.map(map_line)
# 训练模型
model = ALS.train(rdd, 50, 10)

# 模型评估
# 获得所有用户的预测值
ratingForUsers = model.predictAll(rdd.map(lambda r: (r[0], r[1]))).map(lambda r: ((r[0], r[1]), r[2])).join(rdd.map(lambda r: ((r[0], r[1]), r[2])))
metrics = RegressionMetrics(ratingForUsers.map(lambda r: r[1]))
# metrics.meanSquaredError: 均方差

# 保存模型
# model.save(spark, "file:///Users/sonto/Workspace/P1905/spark_example/rec")

# 载入模型
# model = MatrixFactorizationModel.load(spark, "file:///Users/sonto/Workspace/P1905/spark_example/rec")
#
# print(model.predict(889, 381))

# for u in ratingForUsers.collect():
#     print(u)

# rating = model.predict(889, 381)
# print(rating)

